#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库管理脚本
提供常用的数据库操作命令
"""

import asyncio
import sys
import os
import argparse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.database.database import async_engine, AsyncSessionLocal
from src.database.base import Base
from src.models import user  # 导入所有模型
from src.service.user_service import UserService
from src.schemas.user import UserCreate


async def create_tables():
    """创建所有数据库表"""
    print("正在创建数据库表...")
    async with async_engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    print("数据库表创建完成!")


async def drop_tables():
    """删除所有数据库表"""
    print("警告：即将删除所有数据库表!")
    confirm = input("请输入 'yes' 确认删除: ")
    if confirm.lower() == 'yes':
        async with async_engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)
        print("数据库表删除完成!")
    else:
        print("操作已取消")


async def show_tables():
    """显示数据库表信息"""
    print("数据库表信息:")
    print("-" * 50)
    
    for table_name, table in Base.metadata.tables.items():
        print(f"表名: {table_name}")
        print(f"注释: {table.comment or '无'}")
        print("字段:")
        for column in table.columns:
            nullable = "可空" if column.nullable else "非空"
            default = f", 默认值: {column.default.arg if column.default else '无'}"
            print(f"  - {column.name}: {column.type} ({nullable}{default})")
            if hasattr(column, 'comment') and column.comment:
                print(f"    注释: {column.comment}")
        print("-" * 50)


async def count_users():
    """统计用户数量"""
    async with AsyncSessionLocal() as db:
        from sqlalchemy import select, func
        from src.models.user import UserInfo
        
        # 总用户数
        total_stmt = select(func.count(UserInfo.id)).where(UserInfo.del_flag == '0')
        total_result = await db.execute(total_stmt)
        total = total_result.scalar()
        
        # 正常用户数
        active_stmt = select(func.count(UserInfo.id)).where(
            UserInfo.del_flag == '0',
            UserInfo.status == 1
        )
        active_result = await db.execute(active_stmt)
        active = active_result.scalar()
        
        # 禁用用户数
        disabled_stmt = select(func.count(UserInfo.id)).where(
            UserInfo.del_flag == '0',
            UserInfo.status == 0
        )
        disabled_result = await db.execute(disabled_stmt)
        disabled = disabled_result.scalar()
        
        print("用户统计信息:")
        print(f"总用户数: {total}")
        print(f"正常用户: {active}")
        print(f"禁用用户: {disabled}")


async def list_users():
    """列出所有用户"""
    async with AsyncSessionLocal() as db:
        from sqlalchemy import select
        from src.models.user import UserInfo
        
        stmt = select(UserInfo).where(UserInfo.del_flag == '0').order_by(UserInfo.id)
        result = await db.execute(stmt)
        users = result.scalars().all()
        
        print("用户列表:")
        print("-" * 80)
        print(f"{'ID':<5} {'用户名':<15} {'昵称':<15} {'邮箱':<25} {'状态':<8} {'创建时间'}")
        print("-" * 80)
        
        for user in users:
            status = "正常" if user.status == 1 else "禁用"
            create_time = user.create_time.strftime("%Y-%m-%d %H:%M") if user.create_time else "未知"
            print(f"{user.id:<5} {user.username:<15} {user.nick_name or '未设置':<15} {user.email or '未设置':<25} {status:<8} {create_time}")


async def create_admin():
    """创建管理员用户"""
    async with AsyncSessionLocal() as db:
        # 检查是否已存在管理员
        existing_admin = await UserService.get_user_by_username(db, "admin")
        if existing_admin:
            print("管理员用户已存在")
            return
        
        admin_data = UserCreate(
            username="admin",
            password="admin123",
            nick_name="系统管理员",
            email="admin@capcut.com",
            remark="系统默认管理员账户"
        )
        
        try:
            admin_user = await UserService.create_user(db, admin_data)
            print(f"管理员用户创建成功: {admin_user.username}")
        except Exception as e:
            print(f"创建管理员用户失败: {e}")


async def reset_password():
    """重置用户密码"""
    username = input("请输入要重置密码的用户名: ")
    new_password = input("请输入新密码: ")
    
    async with AsyncSessionLocal() as db:
        user = await UserService.get_user_by_username(db, username)
        if not user:
            print("用户不存在")
            return
        
        try:
            # 直接更新密码
            hashed_password = UserService.hash_password(new_password)
            user.password = hashed_password
            await db.commit()
            print(f"用户 {username} 密码重置成功")
        except Exception as e:
            await db.rollback()
            print(f"密码重置失败: {e}")


async def backup_users():
    """备份用户数据到JSON文件"""
    import json
    from datetime import datetime
    
    async with AsyncSessionLocal() as db:
        from sqlalchemy import select
        from src.models.user import UserInfo
        
        stmt = select(UserInfo).where(UserInfo.del_flag == '0')
        result = await db.execute(stmt)
        users = result.scalars().all()
        
        backup_data = {
            "backup_time": datetime.now().isoformat(),
            "total_users": len(users),
            "users": [user.to_dict() for user in users]
        }
        
        filename = f"user_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(backup_data, f, ensure_ascii=False, indent=2)
        
        print(f"用户数据备份完成: {filename}")
        print(f"备份了 {len(users)} 个用户")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='数据库管理工具')
    parser.add_argument('command', choices=[
        'create-tables', 'drop-tables', 'show-tables',
        'count-users', 'list-users', 'create-admin',
        'reset-password', 'backup-users'
    ], help='要执行的命令')
    
    args = parser.parse_args()
    
    try:
        if args.command == 'create-tables':
            await create_tables()
        elif args.command == 'drop-tables':
            await drop_tables()
        elif args.command == 'show-tables':
            await show_tables()
        elif args.command == 'count-users':
            await count_users()
        elif args.command == 'list-users':
            await list_users()
        elif args.command == 'create-admin':
            await create_admin()
        elif args.command == 'reset-password':
            await reset_password()
        elif args.command == 'backup-users':
            await backup_users()
    
    except Exception as e:
        print(f"执行命令时出错: {e}")
    
    finally:
        await async_engine.dispose()


if __name__ == "__main__":
    asyncio.run(main()) 